import numpy as np
import torch
from pyexpat import features
from torch import nn



#1输入字符串
text = "how old are you?"
#2划分数据集
input_seq = []
output_seq = []
window = 5
for i in range(0,len(text)-window,1):
    input_seq.append(text[i:i+window])
    output_seq.append(text[i+window])
    # print(i,input_seq,output_seq)
print(input_seq)
#3进行one-hot编码
chars = set(text) #去重
# print(chars)
chars = sorted(chars)#排序
# print(chars)
#{' ': 0, '?': 1, 'a': 2}
char2int = {char:ind for ind,char in enumerate(chars)}
# print(char2int)
#{0: ' ', 1: '?', 2: 'a'}
chars_dict =dict(enumerate(chars))

print(chars_dict)
#将字符串转成数字编码
input_seq = [[char2int[char]for char in seq] for seq in input_seq]
# print(input_seq)
output_seq = [[char2int[char]for char in seq] for seq in output_seq]

features = np.zeros((len(input_seq),len(chars)),dtype=np.float32)
print(features.shape)
#正式编码
for i,seq in enumerate(input_seq):
    features[i,seq] = 1.0
input_seq = torch.tensor(features,dtype=torch.float32)
features = np.zeros((len(output_seq), len(chars)), dtype=np.float32)
for i,seq in enumerate(output_seq):
    features[i,seq] = 1.0
output_seq = torch.tensor(features,dtype=torch.float32)
print(output_seq)
#4定义前向模型
class Model(nn.Module):
    def __init__(self,input_size,hidden_size,output_size):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(input_size,hidden_size)
        self.fc2 = nn.Linear(hidden_size,output_size)
    def forward(self,x):
        x = nn.functional.relu(self.fc1(x))
        return self.fc2(x)
model = Model(len(chars),32,len(chars))
#5定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.001)
#6开始迭代
epochs = 1000
for epoch in range(1,epochs+1):
    output = model(input_seq)
    loss = loss_fn(output,output_seq)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 7显示频率设置
    if epoch % 2 == 0 or epoch ==0:
        print(f"epoch: {epoch}, loss: {loss}")

#预测下一个字符
input_test = "how o"
input_test = [char2int[char] for char in input_test]
features = np.zeros(len(chars), dtype=np.float32)
for seq in input_test:
    features[seq] = 1.0
input_test = torch.tensor(features,dtype=torch.float32)
out  = model(input_test)
print(out)
print(chars_dict[torch.argmax(out).item()])


